import numpy as np
import matplotlib.pyplot as plt
from scipy import signal
from scipy.fft import fft, fftfreq
import pywt
import warnings

warnings.filterwarnings('ignore')

# 设置全局字体以显示中文（根据你的系统调整字体名称）
plt.rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题


def stft_feature_extraction(signal_data, fs, nperseg=256, noverlap=None, window='hann'):
    """
    使用短时傅里叶变换(STFT)提取时频域特征

    参数:
        signal_data: 一维数组，输入信号数据
        fs: 整数，采样频率
        nperseg: 整数，每个段的长度，默认为256
        noverlap: 整数，段之间重叠的点数，默认为None（使用nperseg//2）
        window: 字符串或数组，指定窗函数，默认为'hann'

    返回:
        stft_features: 字典，包含STFT提取的特征
        stft_data: 字典，包含STFT的原始数据用于可视化
    """
    # 默认重叠为段长的一半
    if noverlap is None:
        noverlap = nperseg // 2

    # 计算STFT
    f, t, Zxx = signal.stft(signal_data, fs=fs, window=window,
                            nperseg=nperseg, noverlap=noverlap)

    # 计算幅度谱
    magnitude = np.abs(Zxx)

    # 预先计算所有需要的值，避免在字典未完全构建前引用
    spectral_centroid = np.sum(f[:, np.newaxis] * magnitude, axis=0) / np.sum(magnitude, axis=0)

    # 使用预先计算的值来计算带宽
    spectral_bandwidth = np.sqrt(
        np.sum((f[:, np.newaxis] - spectral_centroid) ** 2 * magnitude, axis=0) /
        np.sum(magnitude, axis=0)
    )

    # 现在构建特征字典
    stft_features = {
        'spectral_centroid': spectral_centroid,
        'spectral_bandwidth': spectral_bandwidth,
        'spectral_energy': np.sum(magnitude ** 2, axis=0),
        'spectral_entropy': -np.sum(magnitude * np.log2(magnitude + 1e-10), axis=0),
        'time_points': t,
        'freq_points': f
    }

    # 保存原始数据用于可视化
    stft_data = {
        'f': f,
        't': t,
        'Zxx': Zxx,
        'magnitude': magnitude
    }

    return stft_features, stft_data


def wavelet_feature_extraction(signal_data, wavelet='cmor1.0-1.0', scales=None, sampling_period=1.0):
    """
    使用小波变换提取时频域特征

    参数:
        signal_data: 一维数组，输入信号数据
        wavelet: 字符串，小波类型，默认为'cmor1.0-1.0'
        scales: 数组，尺度参数，默认为None（自动生成）
        sampling_period: 浮点数，采样周期，默认为1.0

    返回:
        wavelet_features: 字典，包含小波变换提取的特征
        wavelet_data: 字典，包含小波变换的原始数据用于可视化
    """
    # 自动生成尺度参数（如果未提供）
    if scales is None:
        scales = np.arange(1, 128)

    # 计算连续小波变换(CWT)
    coefficients, frequencies = pywt.cwt(signal_data, scales, wavelet, sampling_period=sampling_period)

    # 计算小波系数幅度
    coeff_magnitude = np.abs(coefficients)

    # 提取小波特征
    wavelet_features = {
        'energy_by_scale': np.sum(coeff_magnitude ** 2, axis=1),
        'max_coeff_by_scale': np.max(coeff_magnitude, axis=1),
        'mean_coeff_by_scale': np.mean(coeff_magnitude, axis=1),
        'total_energy': np.sum(coeff_magnitude ** 2),
        'scales': scales,
        'frequencies': frequencies
    }

    # 保存原始数据用于可视化
    wavelet_data = {
        'coefficients': coefficients,
        'coeff_magnitude': coeff_magnitude,
        'scales': scales,
        'frequencies': frequencies
    }

    return wavelet_features, wavelet_data


def plot_stft_results(stft_data, title="短时傅里叶变换(STFT)谱图"):
    """
    可视化STFT结果

    参数:
        stft_data: 字典，包含STFT的原始数据
        title: 字符串，图表标题
    """
    plt.figure(figsize=(12, 8))

    # 绘制谱图
    plt.subplot(2, 1, 1)
    plt.pcolormesh(stft_data['t'], stft_data['f'], np.abs(stft_data['Zxx']),
                   shading='gouraud', cmap='viridis')
    plt.title(title)
    plt.ylabel('频率 [Hz]')
    plt.xlabel('时间 [秒]')
    plt.colorbar(label='幅度')

    # 绘制几个时间点的频谱切片
    plt.subplot(2, 1, 2)
    time_indices = [0, len(stft_data['t']) // 2, -1]
    colors = ['r', 'g', 'b']
    labels = ['开始', '中间', '结束']

    for i, idx in enumerate(time_indices):
        plt.plot(stft_data['f'], np.abs(stft_data['Zxx'][:, idx]),
                 color=colors[i], label=f'{labels[i]}时刻频谱')

    plt.title('不同时刻的频谱切片')
    plt.ylabel('幅度')
    plt.xlabel('频率 [Hz]')
    plt.legend()
    plt.grid(True)
    plt.tight_layout()
    plt.show()


def plot_wavelet_results(wavelet_data, title="小波变换尺度图"):
    """
    可视化小波变换结果

    参数:
        wavelet_data: 字典，包含小波变换的原始数据
        title: 字符串，图表标题
    """
    plt.figure(figsize=(12, 8))

    # 绘制小波尺度图
    plt.subplot(2, 1, 1)
    plt.imshow(np.abs(wavelet_data['coefficients']),
               extent=[0, 1, wavelet_data['scales'][-1], wavelet_data['scales'][0]],
               aspect='auto', cmap='viridis', interpolation='nearest')
    plt.title(title)
    plt.ylabel('尺度')
    plt.xlabel('时间')
    plt.colorbar(label='系数幅度')

    # 绘制能量随尺度的变化
    plt.subplot(2, 1, 2)
    energy_by_scale = np.sum(np.abs(wavelet_data['coefficients']) ** 2, axis=1)
    plt.plot(wavelet_data['scales'], energy_by_scale)
    plt.title('能量随尺度的变化')
    plt.ylabel('能量')
    plt.xlabel('尺度')
    plt.grid(True)
    plt.tight_layout()
    plt.show()


# 测试代码
if __name__ == "__main__":
    # 生成测试信号（非平稳信号：频率随时间变化）
    fs = 1000  # 采样频率
    t = np.linspace(0, 1, fs)  # 时间向量

    # 创建频率随时间变化的信号（线性调频信号）
    f0 = 5  # 起始频率
    f1 = 50  # 结束频率
    signal_data = np.sin(2 * np.pi * (f0 * t + (f1 - f0) * t ** 2 / 2))

    # 添加一些瞬态冲击
    impulse_times = [0.3, 0.7]
    for impulse_time in impulse_times:
        idx = int(impulse_time * fs)
        signal_data[idx:idx + 50] += 0.5 * np.exp(-np.arange(0, 50) / 10)

    # 添加噪声
    signal_data += 0.1 * np.random.normal(size=len(signal_data))

    # 测试STFT特征提取
    print("进行STFT特征提取...")
    stft_features, stft_data = stft_feature_extraction(signal_data, fs, nperseg=128, noverlap=64)
    print("STFT特征提取完成")
    print(f"频谱重心形状: {stft_features['spectral_centroid'].shape}")
    print(f"频谱带宽形状: {stft_features['spectral_bandwidth'].shape}")

    # 测试小波变换特征提取
    print("\n进行小波变换特征提取...")
    wavelet_features, wavelet_data = wavelet_feature_extraction(signal_data, wavelet='cmor1.0-1.0',
                                                                sampling_period=1 / fs)
    print("小波变换特征提取完成")
    print(f"各尺度能量形状: {wavelet_features['energy_by_scale'].shape}")
    print(f"总能量: {wavelet_features['total_energy']:.4f}")

    # 可视化原始信号
    plt.figure(figsize=(12, 4))
    plt.plot(t, signal_data)
    plt.title('测试信号（非平稳信号：频率随时间变化 + 瞬态冲击 + 噪声）')
    plt.xlabel('时间 [秒]')
    plt.ylabel('幅度')
    plt.grid(True)
    plt.tight_layout()
    plt.show()

    # 可视化STFT结果
    plot_stft_results(stft_data, "测试信号的STFT分析")

    # 可视化小波变换结果
    plot_wavelet_results(wavelet_data, "测试信号的小波变换分析")

    # 打印部分特征值
    print("\n=== STFT特征示例 ===")
    print(f"频谱重心（前5个值）: {stft_features['spectral_centroid'][:5]}")
    print(f"频谱带宽（前5个值）: {stft_features['spectral_bandwidth'][:5]}")

    print("\n=== 小波变换特征示例 ===")
    print(f"最大系数幅度（前5个尺度）: {wavelet_features['max_coeff_by_scale'][:5]}")
    print(f"平均系数幅度（前5个尺度）: {wavelet_features['mean_coeff_by_scale'][:5]}")